import asyncio
import copy
import datetime
import functools
import logging
import pytz
import signal
import typing
from concurrent.futures import ThreadPoolExecutor

from absl import app, flags
from tornado.concurrent import run_on_executor
from tornado.ioloop import IOLoop, PeriodicCallback
from coin.base.datetime_util import to_timestamp_int
from coin.base.logging import init_logging
from coin.support.alarm.app.alarm_stream_processor import define_alarm_flags
from coin.tool.strat_monitor.logic.stream_processor_manager import (
    StrategyStreamProcessorManager)
from coin.tool.strat_monitor.util.monitor_util import (get_active_strategy,
                                                       get_topic_from_strat,
                                                       separate_into_group)
from coin.util.queue.config import KafkaConfig
from coin.util.queue.constants import TopicType


FLAGS = flags.FLAGS


class AlarmProcessorManager(StrategyStreamProcessorManager):
  def __init__(self,
               *,
               ioloop: IOLoop,
               coin_env: str,
               max_workers: int,
               **alarm_args,
    ):
    self._coin_env = coin_env
    self._alarm_args = alarm_args
    strategy_names = alarm_args['strategy_name']
    if strategy_names is not None:
      strategy_names = [elem.strip() for elem in strategy_names.split(',')]
    business_units = alarm_args['business_unit']
    if business_units is not None:
      business_units = [elem.strip() for elem in business_units.split(',')]
    self._strategy_names = strategy_names
    self._business_units = business_units
    self._kafka_config_filename = alarm_args['kafka_config_filename']
    self._kafka_topic_type = TopicType[alarm_args['kafka_topic_type']]
    self._kafka_config = KafkaConfig.from_cmd_config(self._kafka_config_filename)
    self._executor = ThreadPoolExecutor(max_workers=8)
    strat_buckets = separate_into_group(self._get_strategies(), max_workers)
    topic_buckets = [
      get_topic_from_strat(strats, self._kafka_config, topic_type=self._kafka_topic_type)
      for strats in strat_buckets
    ]
    super().__init__(ioloop=ioloop,
                     exit_after_min=alarm_args['exit_after_min'],
                     topic_groups=topic_buckets)

  def gen_sub_monitor_cmd(self, strats: typing.List[str]) -> str:
    manipulated_args = ['strategy_name', 'business_unit']
    for arg in manipulated_args:
      assert arg in self._alarm_args, arg
    alarm_args = copy.deepcopy(self._alarm_args)
    alarm_args['strategy_name'] = ','.join(list(set(strats)))
    del alarm_args['business_unit']
    alarm_args = {key: value for key, value in alarm_args.items() if value is not None}
    cmd = 'source ${HOME}/miniconda3/etc/profile.d/conda.sh' + ' && ' + \
          f'conda activate {self._coin_env}' + ' && ' + \
          './pyrunner coin/support/alarm/app/alarm_stream_processor.py'
    args_str = ' '.join([i for i in flags.flag_dict_to_args(alarm_args)])
    if args_str != '':
      args_str = ' ' + args_str
    cmd += args_str
    return cmd

  def _get_strategies(self) -> typing.List[str]:
    if self._strategy_names is None:
      assert self._business_units is not None
      threshold = to_timestamp_int(
          datetime.datetime.utcnow().replace(tzinfo=pytz.UTC) -
          datetime.timedelta(days=1))
      strats = get_active_strategy(
          business_unit=self._business_units, update_timestamp_threshold=threshold)
    else:
      strats = self._strategy_names
    return strats

  @run_on_executor(executor='_executor')
  def _query_topics(self) -> typing.List[str]:
    strats = self._get_strategies()
    topics = get_topic_from_strat(
        strats, self._kafka_config, topic_type=self._kafka_topic_type)
    return topics

  async def check_new_topics(self):
    logging.info('Check new topics')
    topics = await asyncio.wait_for(self._query_topics(), 10)
    self.handle_topics(topics)


def main(_, *, alarm_flags):
  init_logging()

  ioloop = IOLoop.current()
  flag_dict = flags.FLAGS.flag_values_dict()
  alarm_args = {
      key: flag_dict[key] for key, value in flag_dict.items() if key in alarm_flags
  }
  processor_mgr = AlarmProcessorManager(
      ioloop=ioloop,
      max_workers=FLAGS.max_workers,
      coin_env=FLAGS.coin_env,
      **alarm_args)

  PeriodicCallback(processor_mgr.check_new_topics, 30*1000).start()
  signal.signal(signal.SIGTERM, processor_mgr.signal_handler)
  signal.signal(signal.SIGINT, processor_mgr.signal_handler)
  try:
    ioloop.start()
  except (KeyboardInterrupt, SystemExit):
    return


if __name__ == '__main__':
  logging.basicConfig(level='INFO', format='%(levelname)8s %(asctime)s %(message)s')

  alarm_flags = define_alarm_flags()

  flags.DEFINE_string('coin_env',
                      '~/miniconda3/envs/coin2_motion_env_migration',
                      'coin_env path relative to remote_home')
  flags.DEFINE_integer('max_workers', 24, 'Max number of workers.')

  app.run(functools.partial(main, alarm_flags=alarm_flags))
